from matplotlib.figure import Figure
from dataclasses import dataclass, asdict, field
from matplotlib import pyplot as plt
from matplotlib.gridspec import GridSpec, SubplotSpec
from matplotlib.axes import Axes
from typing import Literal
import numpy as np
import re
import os
from typing import Callable, Concatenate
from cartopy import crs as ccrs
from cartopy.crs import Projection
from rich.progress import Progress
from xarray import DataArray
from matplotlib.transforms import Bbox, BboxTransform

from . import core
from .axes_tools import AxesTools

from .length_units import (
    Length, LengthUnitFactory, UNITS_NAME, as_unit_value, Inch,
    as_inch, as_inch_value, Pt, LenType
)

class BreakMultiLoop(Exception):
    pass

class SubplotParams(AxesTools):
    """循环组图时子图的参数类，包含行数、列数、工具等对象。

    The parameter class of the subplot when looping through the group plot,
    including the number of rows, columns, and tool objects.
    """

    def __init__(self, nrows: int, ncols: int, row: int, col: int,
                    ax: Axes, row_index: int, col_index: int):
        """创建一个 SubplotParams 实例。

        Create an instance of SubplotParams.

        Parameters
        ----------
        nrows : int
            子图的行数

            The number of rows of the subplots
        ncols : int
            子图的列数

            The number of columns of the subplots
        row : int
            子图的行号

            The row number of the subplot
        col : int
            子图的列号

            The column number of the subplot
        ax : Axes
            子图对象

            The Axes object of the subplot
        row_index : int
            子图按行优先的序号

            The serial number of the subplot in row-major order
        col_index : int
            子图按列优先的序号

            The serial number of the subplot in column-major order
        """
        super().__init__(ax)
        self.nrows = nrows
        self.ncols = ncols
        self.row = row
        self.col = col
        self.ax = ax
        self.row_index = row_index
        self.col_index = col_index

        self.is_left = col == 0
        self.is_right = col == ncols - 1
        self.is_top = row == 0
        self.is_bottom = row == nrows - 1


class MoiFigure(Figure):
    """Moisten 的 Figure 类，拓展了 MPL 原来的 Figure 类，提供了更多的功能。

    The Figure subclass of Moist, extends the original Figure class of MPL,
    providing more features.

    Usage
    -----
    首先使用 `plt.figure()` 创建一个画布，并且使用 FigureClass 参数实例化 MoiFigure，
    加上类型注释可以获得更好的代码提示。

    ```python
    fig: MoiFigure = plt.figure(FigureClass=MoiFigure)
    ```

    后续按照正常的 MPL 用法进行绘图即可。
    """

    _used_custom_font = False


    def __init__(self, moi_config: bool=True, **kwargs):
        """创建一个 MoiFigure 实例。

        Create an instance of MoiFigure.

        Parameters
        ----------
        moi_config : bool, optional
            是否使用 Moisten 默认的配置，包括字体、边框粗细等，默认为 True

            Whether to use the default configuration of Moisten, including
            font, border thickness, etc., by default True
        """
        self._use_a4_size = False
        super().__init__(**kwargs)

        if moi_config:
            core.MoiRcParams().apply()
            self.use_std_font()


    def use_a4_size(self, height: LenType=5, padding: LenType=1):
        """设置用于 A4 纸的画幅，即宽为 8.27 英寸减去两边边距。如果长度是数字，则单位为英寸。

        Set the size for A4 paper, i.e. width is 8.27 inches minus two paddings.
        If the length is a number, the unit is inch.

        Parameters
        ----------
        height : LenType, optional
            画幅的高度

            The height of figure, by default 5
        padding : LenType, optional
            A4 排版时的页左右边距，画幅的实际宽度=A4宽度-2*边距

            The left and right margin when A4 layout, the actual width of the
            figure = A4 width - 2 * margin
        """
        height, padding = as_inch_value(self, height, padding)

        self._use_a4_size = True
        self.set_size_inches(8.27 - padding * 2, height)


    @staticmethod
    def _insert_font_family(font_name: str):
        """插入字体到 rcParams['font.family'] 的数组前面，即设置字体"""
        ff = plt.rcParams['font.family']
        if isinstance(ff, str):
            plt.rcParams['font.family'] = [font_name, ff]
        else:
            plt.rcParams['font.family'] = [font_name] + ff


    @staticmethod
    def use_std_font():
        """
        设置标准字体，将使用 assets/fonts 中的字体文件。使用开源的 Roboto 和 MiSans 字体。

        Set the standard font, which will use the font files in assets/fonts.
        Use the open-source Roboto and MiSans fonts.
        """
        from matplotlib import font_manager

        font_dir = os.path.dirname(os.path.realpath(__file__))
        font_dir = os.path.dirname(font_dir) + "/assets/fonts"

        font_files = font_manager.findSystemFonts(fontpaths=font_dir)

        for font_file in font_files:
            font_manager.fontManager.addfont(font_file)

        MoiFigure._insert_font_family('MiSans')
        MoiFigure._insert_font_family('Roboto')


    @staticmethod
    def use_font(font_file_path: str, font_name: str=None,
                 as_math_font: bool=False):
        """
        使用自定义的字体。此字体将会被插入到 rcParams['font.family'] 最前。
        如果不指定 font_name，将会自动识别字体的名称。
        此方法可以在实例化 SuperFigure 之前调用，这样的话 SuperFigure 就不会设置默认的字体。

        Use a custom font. This font will be inserted at the front of
        rcParams['font.family'].
        If `font_name` is not specified, the name of the font will be
        automatically recognized.
        This method can be called before instantiating `SuperFigure`, so that
        `SuperFigure` will not set the default font.

        Parameters
        ----------
        font_file_path : str
            字体文件的路径
            The path to the font file
        font_name : str
            字体的名称
            The name of the font
        """

        from matplotlib import font_manager

        if font_name is None:
            ttf_len = len(font_manager.fontManager.ttflist)
            afm_len = len(font_manager.fontManager.afmlist)

        font_manager.fontManager.addfont(font_file_path)

        if font_name is None:
            if len(font_manager.fontManager.ttflist) > ttf_len:
                font_name = font_manager.fontManager.ttflist[-1].name
            else:
                font_name = font_manager.fontManager.afmlist[-1].name

        MoiFigure._insert_font_family(font_name)
        if as_math_font:
            plt.rcParams['mathtext.default'] = 'regular'

        MoiFigure._used_custom_font = True


    @staticmethod
    def set_math_font(font_name: Literal['dejavusans', 'dejavuserif', 'cm',
                                         'stix', 'stixsans']=None,
                      use_text_font: bool=False):
        """
        设置数学公式显示的字体，如果想使用正常文字的字体作为数学公式的字体，
        可以设置 `use_text_font=True` 。

        Set the font for displaying latex formulas.
        If you want to use the font of normal text as the font of latex formulas,
        you can set `use_text_font=True`.

        Parameters
        ----------
        font_name : Literal['dejavusans', 'dejavuserif', 'cm', 'stix', 'stixsans'], optional
            公式字体名称

            The name of the font for latex formulas
        use_text_font : bool, optional
            使用正常字体作为公式字体, 默认为 False

            Use the font of normal text as the font of latex formulas,
            by default False
        """

        if use_text_font:
            plt.rcParams['mathtext.default'] = 'regular'
        elif font_name is not None:
            plt.rcParams['mathtext.fontset'] = font_name


    def unit(self, name: UNITS_NAME) -> LengthUnitFactory:
        """创建一个长度单位对象，以用于设置边距、大小等。

        Create a length unit object for setting margins, sizes, etc.

        Usage
        -----
        ```python
        pt = fig.unit('pt')  # create a pt unit factory
        fig.subplots_adjust(left=pt(10))  # set left margin to 10pt
        ```

        Parameters
        ----------
        name : UNITS_NAME
            单位的名称，有效值为 'inch', 'pt', 'fs', 'px', 'relx', 'rely', 'rel'

        Returns
        -------
        LengthUnitFactory
            单位工厂对象
        """
        return LengthUnitFactory(self, name)


    def subplots_adjust(self, left:LenType=None, bottom:LenType=None,
                        right:LenType=None, top:LenType=None,
                        wspace:float=None, hspace:float=None) -> None:
        """调整子图的布局参数，可以使用 MoiFigure 的长度单位或者直接使用 float。

        Adjust the subplot layout parameters, you can use the length unit of
        MoiFigure or directly use float.

        Usage
        -----
        ```python
        pt = fig.unit('pt')  # create a pt unit factory
        fig.subplots_adjust(left=pt(10), right=pt(10))  # set layout in pt unit.
        ```

        Parameters
        ----------
        left, bottom, right, top : LenType, optional
            左、下、右、上的边距，可以使用 MoiFigure 的长度单位或者直接使用 float。

            The left, bottom, right, and top margins, you can use the length unit
            of MoiFigure or directly use float.
        wspace, hspace : float, optional
            子图之间的水平和垂直间距，实际间距为各个子图平均的宽高，只能使用 float。

            The horizontal and vertical spacing between subplots, the actual
            spacing is the average width and height of each subplot, can only
            use float.

        """
        left, right = as_unit_value('relx', left, right, fig=self)
        bottom, top = as_unit_value('rely', bottom, top, fig=self)

        if isinstance(wspace, Length) or isinstance(hspace, Length):
            raise ValueError("wspace and hspace must be float, not unit.")

        super().subplots_adjust(left=left, bottom=bottom, right=right,
                                top=top, wspace=wspace, hspace=hspace)


    def extend(self, left: LenType=0, bottom: LenType=0, right: LenType=0, 
               top: LenType=0) -> None:
        """拓展画幅的大小，同时不改变子图的绝对大小。默认以英寸为单位。

        Extend the size of the figure without changing the absolute size of 
        the subplots. The default unit is inch.

        Parameters
        ----------
        left, right, bottom, top : LenType, optional
            画幅的四边的拓展大小，可以使用 MoiFigure 的长度单位或者直接使用 float。

            The extension size of the four sides of the figure, you can use the
            length unit of MoiFigure or directly use float.
        """
        left, bottom, right, top = as_inch_value(self, left, bottom, right, top)

        extend = np.array([[left, bottom], [right, top]])

        origin_width = self.get_figwidth()
        origin_height = self.get_figheight()
        origin_bbox = Bbox.from_extents(0, 0, origin_width, origin_height)

        new_bbox = Bbox.from_extents(0, 0,
            origin_width + extend[:, 0].sum(),
            origin_height + extend[:, 1].sum()
        )

        offset = extend[0] / np.array([new_bbox.width, new_bbox.height])

        trans = BboxTransform(new_bbox, origin_bbox)
        for ax in self.axes:
            points = ax.get_position().get_points()
            points = trans.transform(points)
            points += offset[None, :]
            ax.set_position(
                Bbox.from_extents(*points.flatten())
            )

        self.set_size_inches( new_bbox.width, new_bbox.height)


    def append_colorbar(self, mappable,
                    ax: Axes | np.ndarray[Axes] | list[Axes] | SubplotSpec=None,
                    position: Literal['right', 'bottom']='right',
                    margin: float=Pt(5), width: float=Pt(15),
                    pad: float=0, include_deco: bool=True,
                    re_layout: bool=True, resize_fig: bool=True, **kwargs):
        """附加一个 colorbar 到指定的子图边上，默认情况下，添加colorbar后会重新调整所有子图布局，
        且扩大画幅大小以保持美观。与 `colorbar()` 方法不同，它不会修改原子图的布局，因此更加美观。

        Append a colorbar to the specified side of the subplot. By default, 
        after adding the colorbar, the layout of all subplots will be adjusted,
        and the size of the figure will be enlarged to keep it beautiful.
        Different from the `colorbar()` method, it will not modify the layout of
        the original subplots, so it is more beautiful.

        Parameters
        ----------
        mappable : mappable
            可以应用于 `colorbar()` 的对象

            The object that can be applied to `colorbar()`
        ax : Axes | np.ndarray[Axes] | list[Axes] | SubplotSpec, optional
            指定附加的子图 Axes，若无则自动获取. 可以是一个或一组 Axes 对象，或者是一个 SubplotSpec 对象。

            The specified subplot Axes to append the colorbar,
            if None, it will be automatically obtained.
            It can be a single or a group of Axes objects, or a SubplotSpec object.
        position : Literal['right', 'bottom'], optional
            colorbar 附加的位置

            The position where the colorbar is appended, by default 'right'
        margin : LenType, optional
            colorbar 与子图边框的距离，如果是数值，则单位为相对单位（0-1）

            The distance between the colorbar and the subplot border, 
            by default 5 pt
        width : LenType, optional
            colorbar 的宽度，如果是数值，则单位为相对单位（0-1）

            The width of the colorbar, by default 15 pt
        pad : LenType, optional
            colorbar 头尾边距，相当于 shrink，值越大，colorbar 越短，
            如果是数值，则单位为相对单位（0-1）

            The padding of the colorbar, by default 0 pt
        include_deco : bool, optional
            是否考虑子图的刻度、刻度标签等要素，避免 colorbar 于其重合

            Whether to consider the elements of the subplot, such as ticks and
            tick labels, to avoid overlapping with the colorbar, by default True.
        re_layout : bool, optional
            是否重新布局所有子图，以避免重叠

            Whether to re-layout all subplots to avoid overlapping, by default True
        resize_fig : bool, optional
            是否调整画幅大小，以防止超出画幅。

            Whether to resize the figure to prevent it from exceeding the
            figure size, by default True
        **kwargs
            其他 `colorbar()` 方法的参数

            Other parameters of the `colorbar()` method

        Returns
        -------
        colorbar
            返回 colorbar 对象
            Return the colorbar object
        """

        if ax is None:
            ax = self.gca()

        # 获取子图的位置
        if isinstance(ax, SubplotSpec):
            pos = ax.get_position(self)
            include_deco = False
        else:
            pos = core.bbox_xy(ax, fig=self)

        margin, width, pad = as_unit_value('relx' if position == 'right' else 'rely',
                                           margin, width, pad, fig=self)

        # 如果考虑子图的刻度、刻度标签等要素，使用 tightbbox
        if include_deco:
            tbox = core.bbox_xy(ax, fig=self, tight_box=True)

            if position == 'right':
                margin += tbox.x1 - pos.x1
            elif position == 'bottom':
                margin += pos.y0 - tbox.y0

        if position == 'right':
            left = pos.x1 + margin
            bottom = pos.y0 + pad
            height = pos.height - pad * 2

        elif position == 'bottom':
            left = pos.x0 + pad
            bottom = pos.y0 - margin - width
            height = width
            width = pos.width - pad * 2

        # 保存原图最右边的 axes 距离画幅右边的距离
        if resize_fig:
            if position == 'right':
                # 计算原图最右边的 axes 距离画幅右边的距离
                ax_dis = max([core.bbox_xy(ax, self, True).xmax for ax in self.axes])
                ax_dis = (1 - ax_dis) * self.get_figwidth()
            elif position == 'bottom':
                ax_dis = min([core.bbox_xy(ax, self, True).ymin for ax in self.axes])
                ax_dis = ax_dis * self.get_figheight()


        # 添加 colorbar
        cax = self.add_axes([left, bottom, width, height])
        orientation = 'vertical' if position == 'right' else 'horizontal'
        colorbar =  self.colorbar(mappable, cax=cax, orientation=orientation,
                                  **kwargs)
        cb_pos = colorbar.ax.get_position()
        cb_tight_bbox = core.bbox_xy(colorbar.ax, fig=self, tight_box=True)

        # 重新布局
        if re_layout:
            err = 0.03

            # 搜索有无同行或同列的子图（colorbar），如果有则不需要重新布局了
            need_layout = True
            for ax in self.axes:
                if ax is colorbar.ax:
                    continue
                ax_pos = ax.get_position()
                if position == 'right':
                    if ax_pos.xmin >= cb_pos.xmin - err and \
                        ax_pos.xmax <= cb_pos.xmax + err:
                        need_layout = False
                        break
                elif position == 'bottom':
                    if ax_pos.ymin >= cb_pos.ymin - err and \
                        ax_pos.ymax <= cb_pos.ymax + err:
                        need_layout = False
                        break

            if need_layout:

                for ax in self.axes:
                    if ax is colorbar.ax:
                        continue
                    ax_pos = ax.get_position()
                    if position == 'right':
                        w = cb_tight_bbox.width + margin
                        if ax_pos.xmin < cb_pos.xmin - err and \
                            ax_pos.xmax > cb_pos.xmax + err:
                            ax.set_position(Bbox.from_extents(
                                ax_pos.x0,
                                ax_pos.y0,
                                ax_pos.x1 + w,
                                ax_pos.y1
                            ))
                        elif ax_pos.xmin >= cb_pos.xmin - err:
                            ax.set_position(Bbox.from_extents(
                                ax_pos.x0 + w,
                                ax_pos.y0,
                                ax_pos.x1 + w,
                                ax_pos.y1
                            ))
                    elif position == 'bottom':
                        w = cb_tight_bbox.height + margin
                        if ax_pos.ymin < cb_pos.ymin - err and \
                            ax_pos.ymax > cb_pos.ymax + err:
                            ax.set_position(Bbox.from_extents(
                                ax_pos.x0,
                                ax_pos.y0 - w,
                                ax_pos.x1,
                                ax_pos.y1
                            ))
                        elif ax_pos.ymax <= cb_pos.ymax + err:
                            ax.set_position(Bbox.from_extents(
                                ax_pos.x0,
                                ax_pos.y0 - w,
                                ax_pos.x1,
                                ax_pos.y1 - w
                            ))

        # 调整画幅大小
        if resize_fig:
            if position == 'right':
                ax_right = max([core.bbox_xy(ax, self, True).xmax for ax in self.axes])
                ax_right = (1 - ax_right) * self.get_figwidth()
                self.extend(right=ax_dis - ax_right)
            elif position == 'bottom':
                ax_bottom = min([core.bbox_xy(ax, self, True).ymin for ax in self.axes])
                ax_bottom = ax_bottom * self.get_figheight()
                self.extend(bottom=ax_dis - ax_bottom)

        return colorbar


    def ratio_gridspec(self, nrows: int, ncols: int,
            ratio: float|tuple[float, float]|tuple[float,float,float,float],
            subplot_width: LenType=None, gridspec_width: LenType=None,
            left_margin: LenType=Pt(40), right_margin: LenType=Pt(20),
            top_margin: LenType=Pt(20), bottom_margin: LenType=Pt(20),
            hspace: LenType=Pt(20), wspace: LenType=Pt(20),
            extend_fig: bool=True) -> GridSpec:
        """规则固定比例的子图布局（邮票图），组图整体的位置将以 `left_margin`、`top_margin`
        相对于画幅左上角来定位，如果组图的大小超出画幅，将会根据四边的 `margin` 来调整画幅大小。

        需要使用 `subplot_width` 或 `gridspec_width` 来指定子图的宽度。

        Regular fixed ratio subplot layout (postage stamp plot), the overall
        position of the group plot will be located relative to the upper left
        corner of the figure by `left_margin` and `top_margin`.
        If the size of the group plot exceeds the figure, the figure size will
        be adjusted according to the `margin` on the four sides.

        You need to use `subplot_width` or `gridspec_width` to specify the width
        of the subplot.

        Parameters
        ----------
        nrows, ncols : int
            子图的行数和列数

            The number of rows and columns of the subplots
        ratio : float|tuple[float, float]|tuple[float,float,float,float]
            子图的宽高比例，可以是一个浮点数，或者一个包含两个元素的元组，或者一个包含四个元素的元组。

            The width-to-height ratio of the subplots, which can be a floating
            point number, or a tuple containing two elements, or a tuple
            containing four elements.
        subplot_width : LenType, optional
            子图的宽度，如果是数值则单位为英寸。

            The width of the subplots, if it is a number, the unit is inch.
        gridspec_width : LenType, optional
            gridspec 组图整体的宽度，如果是数值则单位为英寸。

            The width of the gridspec group plot as a whole, if it is a number,
            the unit is inch.
        left_margin, right_margin, top_margin, bottom_margin : LenType, optional
            画幅的四边边距，如果是数值则单位为英寸。

            The margin of the four sides of the figure, if it is a number, the
            unit is inch.
        hspace, wspace : LenType, optional
            子图之间的水平和垂直间距，如果是数值则单位为英寸。

            The horizontal and vertical spacing between subplots, if it is a
            number, the unit is inch.
        extend_fig : bool, optional
            如果画幅大小超出，是否调整画幅大小

            If the size of the figure exceeds, whether to adjust the size of the figure.
        """

        # 获取子图比例
        if not isinstance(ratio, float) and not isinstance(ratio, int):
            if len(ratio) == 2:
                ratio = ratio[0] / ratio[1]
            elif len(ratio) == 4:
                ratio = (ratio[1] - ratio[0]) / (ratio[3] - ratio[2])
            else:
                raise ValueError("The length of ratio must be 2 or 4.")

        if subplot_width is None and gridspec_width is None:
            raise ValueError("You must specify either subplot_width or gridspec_width.")

        if subplot_width is not None and gridspec_width is not None:
            raise ValueError("You can't specify both subplot_width and gridspec_width.")

        wspace, hspace = as_inch(self, wspace, hspace)
        left_margin, right_margin, top_margin, bottom_margin = \
            as_inch(self, left_margin, right_margin, top_margin, bottom_margin)

        # 计算子图、gridspec、画幅大小
        if subplot_width is not None:
            subplot_width = as_inch(self, subplot_width)
            subplot_height = subplot_width / ratio
            gridspec_width = subplot_width * ncols + wspace * (ncols - 1)

        elif gridspec_width is not None:
            gridspec_width = as_inch(self, gridspec_width)
            subplot_width = (gridspec_width - wspace * (ncols - 1)) / ncols
            subplot_height = subplot_width / ratio

        gridspec_height = subplot_height * nrows + hspace * (nrows - 1)
        fig_width = gridspec_width + left_margin + right_margin
        fig_height = gridspec_height + top_margin + bottom_margin

        # 如果画幅大小超出，调整画幅大小
        if extend_fig:
            dx, dy = 0, 0
            if fig_width.to_inch() > self.get_figwidth():
                dx = fig_width.to_inch() - self.get_figwidth()
            if fig_height.to_inch() > self.get_figheight():
                dy = fig_height.to_inch() - self.get_figheight()
            if dx + dy > 0:
                self.set_size_inches(self.get_figwidth() + dx,
                                    self.get_figheight() + dy)

        gs = GridSpec(nrows, ncols, self,
                      left=left_margin.to_relx(),
                      top=1 - top_margin.to_rely(),
                      right=(left_margin+gridspec_width).to_relx(),
                      bottom=1 - (top_margin + gridspec_height).to_rely(),
                      hspace=hspace/subplot_height,
                      wspace=wspace/subplot_width)

        return gs


    def add_axes(self, *args, **kwargs):
        # 添加单位支持
        if kwargs.get('rect') is not None:
            rect = kwargs.pop('rect')
        elif len(args) > 0 and (isinstance(args[0], tuple) or\
                                isinstance(args[0], list)):
            rect = args[0]
        else:
            return super().add_axes(*args, **kwargs)

        new_rect = []
        for i, v in enumerate(rect):
            if isinstance(v, Length):
                new_rect.append(v.to_rel('x' if i % 2 == 0 else 'y'))
            else:
                new_rect.append(v)
        return super().add_axes(new_rect, *args[1:], **kwargs)


    def gridspec_by_size(self, nrows: int, ncols: int, width: LenType | list[LenType],
            height: LenType | list[LenType], left_margin: LenType=Pt(40), 
            right_margin: LenType=Pt(20), top_margin: LenType=Pt(20), 
            bottom_margin: LenType=Pt(20), wspace: LenType=Pt(20),
            hspace: LenType=Pt(20), extend_fig: bool=True) -> GridSpec:
        """根据行、列的实际大小来创建 gridspec。组图的左上角根据 `left_margin`、`top_margin`
        来定位，如果组图的大小超出画幅且 `extend_fig=True` ，将会根据四边的 `margin` 来调整
        画幅大小。

        Create gridspec according to the real width and height of rows and columns.
        The upper left corner of the group plot is located according to `left_margin`
        and `top_margin`. If the size of the group plot exceeds the figure and
        `extend_fig=True`, the figure size will be adjusted according to the `margin`
        on the four sides.

        Parameters
        ----------
        nrows : int
            行数  rows
        ncols : int
            列数  cols
        width, height : LenType | list[LenType]
            每列的宽度与每行的高度。如果是单值，则所有行/列相同；如果是列表，则每列/行的宽/高度分别
            为列表中的值。

            The width of each column and the height of each rows.
            If it is a single value, that applies to all rows/columns;
            if it is a list, the width/height of each column/row is the value
            in the list respectively.
        left_margin, right_margin, top_margin, bottom_margin : LenType, optional
            画幅的四边边距，如果是数值则单位为英寸。

            The margin of the four sides of the figure, if it is a number, the
            unit is inch.
        wspace, hspace : LenType, optional
            子图之间的水平和垂直间距，如果是数值则单位为英寸。

            The horizontal and vertical spacing between subplots, if it is a
            number, the unit is inch.
        extend_fig : bool, optional
            如果 Gridspec 大小超出画幅，是否调整画幅大小，默认为 True。

            If the size of the Gridspec exceeds the figure, whether to adjust
            the size of the figure, by default True.

        Returns
        -------
        GridSpec
            返回 GridSpec 对象

            return GridSpec object
        """

        # 根据子图大小创建 gridspec
        if isinstance(width, list) and len(width) == ncols:
            width = [as_inch(self, w) for w in width]
        elif isinstance(width, float) or isinstance(width, int):
            width = [as_inch(self, width)] * ncols
        else:
            raise ValueError("The length of width must be equal to cols.")

        if isinstance(height, list) and len(height) == nrows:
            height = [as_inch(self, h) for h in height]
        elif isinstance(height, float) or isinstance(height, int):
            height = [as_inch(self, height)] * nrows
        else:
            raise ValueError("The length of height must be equal to rows.")

        left_margin, right_margin, top_margin, bottom_margin = \
            as_inch(self, left_margin, right_margin, top_margin, bottom_margin)

        wspace, hspace = as_inch(self, wspace, hspace)
        grid_width = sum(width) + wspace * (ncols - 1)
        grid_height = sum(height) + hspace * (nrows - 1)

        width_ratio = [w.to_inch() for w in width]
        width_mean = np.mean(width_ratio)
        height_ratio = [h.to_inch() for h in height]
        height_mean = np.mean(height_ratio)

        fig_width = left_margin + grid_width + right_margin
        fig_height = top_margin + grid_height + bottom_margin

        # 如果画幅大小超出，调整画幅大小
        if extend_fig:
            dx, dy = 0, 0
            if fig_width.to_inch() > self.get_figwidth():
                dx = fig_width.to_inch() - self.get_figwidth()
            if fig_height.to_inch() > self.get_figheight():
                dy = fig_height.to_inch() - self.get_figheight()
            if dx + dy > 0:
                self.set_size_inches(self.get_figwidth() + dx,
                                    self.get_figheight() + dy)

        gs = GridSpec(nrows, ncols, self, left=left_margin.to_relx(),
                      top=1 - top_margin.to_rely(),
                      right=(left_margin + grid_width).to_relx(),
                      bottom=1 - (top_margin + grid_height).to_rely(),
                      hspace=hspace.to_inch()/height_mean, 
                      wspace=wspace.to_inch()/width_mean,
                      width_ratios=width_ratio, height_ratios=height_ratio)

        return gs


    def _progress_bar(self, title: str, show_progress: bool=True):
        from rich.progress import (Progress, TextColumn, BarColumn, 
        TimeRemainingColumn, MofNCompleteColumn, SpinnerColumn)
        progress = Progress(
            TextColumn(title),
            SpinnerColumn('moon'),
            BarColumn(),
            MofNCompleteColumn(),
            TimeRemainingColumn(True, True),
            disable=not show_progress)
        return progress

    def loop_gridspec(self, gs: GridSpec, 
                      func: Callable[Concatenate[SubplotParams, ...], None],
                      show_progress: bool=True, break_at: int = None,
                      func_args={}, 
                      func_subplot_kwargs:Callable[[int, int], dict]=None, **kwargs):
        
        prog = self._progress_bar("Setting subplots", show_progress)
        prog.start()
        total = (break_at+1) if break_at is not None else gs.nrows * gs.ncols
        task = prog.add_task("Drawing subplots...", total=total)
        results = []

        try:
            for i in range(gs.nrows):
                for j in range(gs.ncols):
                    if func_subplot_kwargs is not None:
                        ax_kwargs = func_subplot_kwargs(i, j)
                    else:
                        ax_kwargs = {}
                    ax_kwargs.update(kwargs)
                    ax = self.add_subplot(gs[i, j], **ax_kwargs)
                    row_index = i * gs.ncols + j
                    col_index = j * gs.nrows + i
                    res = func(
                        SubplotParams(
                            nrows=gs.nrows, ncols=gs.ncols,
                            row=i, col=j,
                            ax=ax,
                            row_index=row_index, col_index=col_index
                        ), **func_args
                    )
                    results.append(res)
                    prog.update(task, advance=1)

                    if break_at is not None and row_index >= break_at:
                        raise BreakMultiLoop

        except BreakMultiLoop:
            pass
        finally:
            prog.stop()
        
        return results


    def fig_title(self, title: str, top_margin: float=Pt(8), 
                  va='top', **kwargs):
        """添加一个总标题到画幅的上方，通过 `top_margin` 来设置标题距离画幅上边框
        的距离。

        Add a main title to the top of the figure. The distance between the 
        title and the top border of the figure is set by `top_margin`.

        Parameters
        ----------
        title : str
            标题内容  title content
        top_margin : float, optional
            到画幅顶部距离, by default Pt(8)

            The distance to the top of the figure, by default Pt(8)
        va : str, optional
            垂直对齐方式，默认为 'top'

            The vertical alignment, by default 'top'

        **kwargs
            其他参数，传递给 `Figure.suptitle()` 方法

            Other parameters, passed to the `Figure.suptitle()` method
        """
        top_margin = as_inch(self, top_margin)
        y = 1 - top_margin.to_rely()
        self.suptitle(title, y=y, va=va, **kwargs)

    def fig_hline(self):
        pass